package class10;

import java.util.PriorityQueue;

/**
 * https://leetcode.cn/problems/find-median-from-data-stream/
 * 数据流的中位数
 * 解题思路：
 * 步骤：创建大根堆和小根堆解决
 * * 操作1：第一个数 5：如果 <= 大根堆顶，进大根堆
 * *                 如果 >= 大根堆顶，进小根堆
 * *
 * * 操作2：大根堆存较小的一半，小根堆存较大的一半
 * *
 */
public class Code13_MedianFinder {

    public static void main(String[] args) {
        MedianFinder medianFinder = new MedianFinder();
        medianFinder.addNum(1);
        medianFinder.addNum(2);
        System.out.println("medianFinder.findMedian() = " + medianFinder.findMedian());
        medianFinder.addNum(3);
        System.out.println("medianFinder.findMedian() = " + medianFinder.findMedian());

    }

    private static class MedianFinder {
        PriorityQueue<Integer> maxHeap; // 大根堆
        PriorityQueue<Integer> minHeap; // 小根堆

        public MedianFinder() {
            maxHeap = new PriorityQueue<>((a, b) -> b - a);
            minHeap = new PriorityQueue<>((a, b) -> a - b);
        }

        public void addNum(int num) {
            if (maxHeap.isEmpty() || num <= maxHeap.peek()) {
                maxHeap.offer(num);
            } else {
                minHeap.offer(num);
            }
            balance();
        }

        /**
         * 1.如果差值为2，则平衡一下，大小根堆
         * 2.数量较大的堆弹出一个，放入另一个堆里
         */
        private void balance() {
            int m = maxHeap.size();
            int n = minHeap.size();
            if (Math.abs(m - n) == 2) {
                if (m > n) {
                    minHeap.offer(maxHeap.poll());
                } else {
                    maxHeap.offer(minHeap.poll());
                }
            }
        }

        public double findMedian() {
            if (maxHeap.size() == minHeap.size()) {
                return (maxHeap.peek() + minHeap.peek()) / 2.0;
            } else {
                return (maxHeap.size() < minHeap.size()) ? minHeap.peek() : maxHeap.peek();
            }
        }
    }
}
